#!/usr/bin/env python3

import rospy
import rosbag
import rospkg
import os
import os.path as osp
from std_msgs.msg import Header
from geometry_msgs.msg import Vector3
from geometry_msgs.msg import Quaternion
from bev_det_bagwriter.msg import CustomizedTfs
from cv_bridge import CvBridge
import cv2
import json
import pickle
import yaml
import numpy as np
from tqdm import tqdm
import concurrent.futures


class BagWriter():
    def __init__(self, cfg_path = None, bag_path = r"bags"):
        self._ros_paths = rospkg.get_ros_paths()
        self.pkg_path = None
        for i in self._ros_paths:
            if "catkin_ws" in i:
                self.pkg_path = osp.join(i, "bev_det_bagwriter")
                break
        assert self.pkg_path, "cannot find this package which name is bev_det_bagwriter, source your workspace"
        if cfg_path is None:
            with open(osp.join(self.pkg_path, "cfg", "config.yaml")) as cfg:
                data_meta = yaml.safe_load(cfg)
                cfg_path = data_meta["database_path"]
                self.select_scene_id = data_meta["select_scene_id"]
        self.cfg_path = osp.join(self.pkg_path, cfg_path)
        self.bag_path = osp.join(self.pkg_path, bag_path)
        self.sensor_channels = list(self.Config.keys())
        self.sensor_channels.remove("scene_token_list")
        self.scene_token_list = self.Config["scene_token_list"]
        self.scene_meta = list(zip(self.select_scene_id, self.scene_token_list))

    @property
    def Config(self):
        for file in os.listdir(self.cfg_path):
            if file.endswith("json"):
                with open(osp.join(self.cfg_path, file), 'r') as config:
                    return json.load(config)
            elif file.endswith("pkl"):
                with open(osp.join(self.cfg_path, file), 'rb') as config:
                    return pickle.load(config)

    def nda2list(self, mat: np.ndarray) -> list:
        return [i for j in mat.tolist() for i in j]

    def quat_nu2ros(self, quat):
        quat[0], quat[1], quat[2], quat[3] = quat[1], quat[2], quat[3], quat[0]
        return quat

    # TODO: construct all the bags here
    def write_bag(self):
        with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
            executor.map(self.write, self.scene_token_list)

    # TODO: construct one bag here
    def write(self, scene_token):
        bridge = CvBridge()
        with rosbag.Bag(osp.join(self.bag_path, scene_token+".bag"), 'w') as outbag:
            for cam in self.sensor_channels:
                image_path = self.Config[cam][scene_token]["path"]
                timestamp = self.Config[cam][scene_token]["timestamp"]
                transformations = self.Config[cam][scene_token]["infos"]
                for img_file, tsp, tfs in tqdm(zip(image_path, timestamp, transformations)):
                    img_cv = cv2.imread(osp.join(self.cfg_path, img_file))
                    secs = tsp / 1e6
                    tsp_msg = Header(stamp=rospy.Time.from_sec(secs), frame_id=cam)
                    img_msg = bridge.cv2_to_imgmsg(cvim=img_cv, header=tsp_msg)

                    cutfs = CustomizedTfs()
                    cutfs.header = tsp_msg
                    cutfs.sensor2ego_translation = Vector3(*(tfs["sensor2ego_translation"].tolist()))
                    cutfs.sensor2ego_rotation = Quaternion(*self.quat_nu2ros((tfs["sensor2ego_rotation"].tolist())))
                    cutfs.ego2global_translation = Vector3(*(tfs["ego2global_translation"].tolist()))
                    cutfs.ego2global_rotation = Quaternion(*self.quat_nu2ros((tfs["ego2global_rotation"].tolist())))
                    cutfs.cam_intrinsic = self.nda2list(tfs["cam_intrinsic"])
                    # print(cutfs)
                    outbag.write(topic=cam, msg=img_msg, t=rospy.Time.from_sec(secs))
                    outbag.write(topic=cam+"_info", msg=cutfs, t=rospy.Time.from_sec(secs))


def main():

    bw = BagWriter()
    bw.write_bag()


if __name__ == "__main__":
    main()
